# python -B -m vehicle_flow_counting.main

from vehicle_flow_counting import EDGE_MARGIN, SHOW_EDGE, INTERVAL_FRAME, DISPLAY_INPUT, DISPLAY_SCALE_FACTOR, DISPLAY_FLAG, DISPLAY_SAVE_FLAG, DISPLAY_PUSH, MODEL_FILE, CLASS_NAMES
from .utils.frame_display import show_frame_scaled
from .detector import YoloTracker
from .counter import FlowCounter
from .utils.video_io import VideoIO

def main():
    # 测试
    # video_path = "../data/25s.mp4"
    # video_path = "rtmp://47.111.102.121/live/djl/"
    # model_path = "../model/yolo11m-obb-visdrone-0.pt"
    # model_path = "../model/yolo11n.pt"

    video_io = VideoIO()
    width, height, _, input_url, classes = video_io.get_params()
    if DISPLAY_INPUT: # 测试显示输入视频流是否正常（正式使用时为 False）
        video_io.play_input_video(input_url)
        return

    tracker = YoloTracker(MODEL_FILE) # 初始化 YOLO tracker

    flow_counter = FlowCounter((height, width), CLASS_NAMES, EDGE_MARGIN, SHOW_EDGE) # 用户可设定阈值 edge_margin 和 筛选类别 class_names
    frame_id = 0
    # 遍历每帧
    for r in tracker.track(input_url, classes=classes, interval_frame = INTERVAL_FRAME): # interval_frame > 1 可减少目标消失重现后导致的id变化
        im0 = r.orig_img.copy()
        # print(im0.shape if im0 is not None else "im0 is None")

        # 获取当前帧所有track (track_id, (cx,cy))
        current_tracks = tracker.parse_tracks(r)

        # 更新统计
        flow_counter.update(current_tracks)
        # 绘制统计数据
        flow_counter.draw_counts(im0)
        flow_counter.get_stats_json()
        # 绘制中心点和ID
        flow_counter.draw_target(im0, current_tracks)

        # 显示(上述统计绘制后的画面，非yolo生成的结果)
        exit_flag = show_frame_scaled(im0, window_name="Vehicle Flow Counting", mode="video", scale_factor=DISPLAY_SCALE_FACTOR, show=DISPLAY_FLAG)
        # 保存上述视频 (原尺寸和帧率)
        if DISPLAY_SAVE_FLAG:
            video_io.write_frame(im0)
        # 推流 (含分辨率调整)
        frame_id = frame_id + 1
        if DISPLAY_PUSH:
            video_io.push_frame(im0)
            print("frame id: ", frame_id)

        if exit_flag:  # 空格键退出
            break

    video_io.release() # 释放资源

if __name__ == "__main__":
    main()
